#!groovy
// -*- mode: groovy -*-

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

AARCH64_IMAGE = "cpu_aarch64:${params.TLCPACK_AARCH64_TAG}"


def cleanup_docker_image(image_type) {
    // Delete images created/tagged by previous steps.
    sh (
        script: "docker rmi -f tlcpack/package-${image_type}",
        label: "Cleanup ${image_type} image",
    )
}


def clone_tlcpack() {
    // clone tlcpack repo directly into Jenkins workspace
    sh (
        script: "git clone ${params.TLCPACK_GIT_REPO_URL} .",
        label: "Clone tlcpack into Jenkins workspace",
    )
}


def clone_tvm() {
    sh (
        script: "git clone --branch ${params.GIT_REF} --depth 1 ${params.TVM_GIT_REPO_URL} --recursive",
        label: "Clone TVM in to /tvm directory",
    )
}


def pull_docker_image(image_type) {
    // The tag is expected to be passed as a Jenkins parameter
    sh (
        script: "docker pull tlcpack/package-${image_type}",
        label: "Pull 'package-${image_type}' from Docker Hub",
    )
}


def sync_package() {
    sh (
        script: """
            python3 common/sync_package.py \
              --cuda ${env.GPU} \
              --package-name ${params.PACKAGE_NAME} \
              --use-public-version \
              ${params.PACKAGE}
        """,
        label: "Sync tlcpack package",
    )
}


def build_wheels(image_type) {
    sh (
        script: """
            ./docker/bash.sh --no-gpu tlcpack/package-${image_type} \
                ./wheel/build_wheel_manylinux.sh --cuda ${env.GPU}
        """,
        label: "Build wheels on ${image_type}",
    )
}


def list_wheels() {
    sh (
        script: "find . -type f -name *.whl",
        label: "Find and list wheels",
    )
}


def test_wheels(image_type) {
    sh (
        script: """
            ./docker/bash.sh --no-gpu tlcpack/package-${image_type} \
                ./wheel/run_tests.sh
        """,
        label: "Test wheels for ${image_type}",
    )
}


def upload_to_pypi() {
    sh (
        script: "python3 -m pip install twine",
        label: "Install python package 'twine'",
    )
    sh (
        script: "python3 -m twine upload tvm/python/repaired_wheels/*",
        label: "Twine upload the repaired wheels to PyPI",
    )
}



pipeline {

    agent { node { label 'ARM-SPOT' } }

    environment {
        GPU = "none"
        DISCORD_WEBHOOK = credentials('discord-webhook-url')
        TWINE_NON_INTERACTIVE = 1
        TWINE_REPOSITORY = "pypi"
        TWINE_USERNAME = "__token__"
        TWINE_PASSWORD = credentials('pypi-api-token')
    }

    options {
        timeout(time: 1, unit: "HOURS")
    }

    parameters {
        string(
            name: "GIT_REF",
            defaultValue: "v0.12.0",
            description: "The branch to checkout and package")
        string(
            name: "PACKAGE",
            defaultValue: "pypi",
            description: "What type of package to create")
        string(
            name: "PACKAGE_NAME",
            defaultValue: "apache-tvm",
            description: "What name the package will be upload to PyPI with")
        string(
            name: "TLCPACK_GIT_REPO_URL",
            defaultValue: "https://github.com/tlc-pack/tlcpack",
            description: "URL for the tlcpack repository")
        string(
            name: "TLCPACK_AARCH64_TAG",
            defaultValue: "2ac0ed7b8",
            description: "Tag for the Docker tlcpack AArch64 image")
        string(
            name: "TVM_GIT_REV",
            defaultValue: "main",
            description: "Git revision to checkout.")
        string(
            name: "TVM_GIT_REPO_URL",
            defaultValue: "https://github.com/apache/tvm",
            description: "URL for the TVM repository")
    }

    stages {

        stage("Package AArch64") {
            steps {
                cleanWs()
                clone_tlcpack()
                clone_tvm()
                dir("tvm") {
                    script {
                        TVM_CURRENT_SHORT_REF = sh (
                            script: "git rev-parse --short HEAD",
                            label: "Set current build display name",
                            returnStdout: true
                            ).trim()
                        currentBuild.displayName = "tlcpack-${TVM_CURRENT_SHORT_REF}"
                    } // script
                } // dir(tvm)
                sync_package()
                pull_docker_image("${AARCH64_IMAGE}")
                build_wheels("${AARCH64_IMAGE}")
                cleanup_docker_image("${AARCH64_IMAGE}")
                list_wheels()
                upload_to_pypi()
            }
        } // stage(Package AArch64)

    } // stages

    post {
        success {
            discordSend (
                description: "New packages published on PyPI for ${AARCH64_IMAGE}.",
                link: "https://pypi.org/project/apache-tvm/#history",
                result: currentBuild.currentResult,
                title: "${JOB_NAME}",
                webhookURL: "${env.DISCORD_WEBHOOK}"
            )
        }
        unsuccessful {
            discordSend (
                description: "Failed to generate tlcpack packages using TVM hash `${TVM_CURRENT_SHORT_REF}`. See logs at ${BUILD_URL}.",
                link: "${BUILD_URL}",
                result: currentBuild.currentResult,
                title: "${JOB_NAME}",
                webhookURL: "${env.DISCORD_WEBHOOK}"
            )
        }
        cleanup {
            cleanWs();
        }
    } // post
} // pipeline
